import tenseal as ts
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from tkinter import *
from tkinter import ttk, messagebox
from typing import List, Union, Tuple
import os

# 数据文件路径配置
DATA_PATH = r"E:\2024-2025第二学期\密码学进展\期末论文"
PRODUCTS_FILE = os.path.join(DATA_PATH, "商品表.xlsx")
PURCHASES_FILE = os.path.join(DATA_PATH, "购买记录表.xlsx")

class PrivacyPreservingSpendingAnalyzer:
    def __init__(self, root):
        self.root = root
        self.root.title("隐私保护消费分析工具")
        self.root.geometry("900x700")

        # 初始化FHE引擎
        self.fhe = FHEEncryption(scheme="CKKS")

        # 数据存储
        self.products_df = None
        self.purchases_df = None
        # 新增：存储明文和加密计算结果
        self.plaintext_results = None
        self.encrypted_results = None

        # 创建UI
        self.create_widgets()

    def create_widgets(self):
        # 数据加载区
        frame_load = LabelFrame(self.root, text="数据加载", padx=10, pady=10)
        frame_load.pack(pady=10, padx=10, fill="x")
        Button(frame_load, text="加载商品表", command=self.load_products).pack(side="left", padx=5)
        Button(frame_load, text="加载购买记录", command=self.load_purchases).pack(side="left", padx=5)

        # 加密参数区
        frame_params = LabelFrame(self.root, text="加密参数", padx=10, pady=10)
        frame_params.pack(pady=10, padx=10, fill="x")
        Label(frame_params, text="多项式模数:").grid(row=0, column=0, sticky="e")
        self.poly_mod_entry = Entry(frame_params)
        self.poly_mod_entry.insert(0, "8192")
        self.poly_mod_entry.grid(row=0, column=1, sticky="w")
        Label(frame_params, text="系数模数位大小:").grid(row=1, column=0, sticky="e")
        self.coeff_mod_entry = Entry(frame_params)
        self.coeff_mod_entry.insert(0, "60,40,40,60")
        self.coeff_mod_entry.grid(row=1, column=1, sticky="w")
        Label(frame_params, text="缩放因子:").grid(row=2, column=0, sticky="e")
        self.scale_entry = Entry(frame_params)
        self.scale_entry.insert(0, "1099511627776")  # 2^40
        self.scale_entry.grid(row=2, column=1, sticky="w")

        # 分析操作区
        frame_analysis = LabelFrame(self.root, text="消费分析", padx=10, pady=10)
        frame_analysis.pack(pady=10, padx=10, fill="x")
        Button(frame_analysis, text="计算客户消费(明文)", command=self.analyze_plaintext).pack(side="left", padx=5)
        Button(frame_analysis, text="计算客户消费(加密)", command=self.analyze_encrypted).pack(side="left", padx=5)

        # 结果展示区
        frame_results = LabelFrame(self.root, text="分析结果", padx=10, pady=10)
        frame_results.pack(pady=10, padx=10, fill="both", expand=True)
        self.tree = ttk.Treeview(frame_results, columns=("CustomerID", "TotalSpent"), show="headings")
        self.tree.heading("CustomerID", text="客户编号")
        self.tree.heading("TotalSpent", text="消费金额")
        self.tree.pack(fill="both", expand=True, padx=5, pady=5)
        Button(frame_results, text="可视化结果", command=self.visualize_results).pack(pady=5)

        # 状态栏
        self.status = Label(self.root, text="就绪", bd=1, relief="sunken", anchor="w")
        self.status.pack(side="bottom", fill="x")

    def update_status(self, message):
        self.status.config(text=message)
        self.root.update()

    def load_products(self):
        try:
            self.products_df = pd.read_excel(PRODUCTS_FILE)
            self.update_status("商品数据已加载")
        except Exception as e:
            messagebox.showerror("错误", f"加载商品数据失败: {str(e)}")

    def load_purchases(self):
        try:
            self.purchases_df = pd.read_excel(PURCHASES_FILE)
            self.update_status("购买记录已加载")
        except Exception as e:
            messagebox.showerror("错误", f"加载购买记录失败: {str(e)}")

    def update_encryption_params(self):
        try:
            poly_mod = int(self.poly_mod_entry.get())
            coeff_mod = [int(x) for x in self.coeff_mod_entry.get().split(",")]
            scale = int(self.scale_entry.get())
            self.fhe = FHEEncryption(
                scheme="CKKS",
                poly_modulus_degree=poly_mod,
                coeff_mod_bit_sizes=coeff_mod,
                scale=scale
            )
            self.update_status("加密参数已更新")
        except Exception as e:
            messagebox.showerror("错误", f"更新加密参数失败: {str(e)}")

    def analyze_plaintext(self):
        if self.products_df is None or self.purchases_df is None:
            messagebox.showerror("错误", "请先加载商品数据和购买记录")
            return
        try:
            self.update_status("正在计算明文消费数据...")
            merged = pd.merge(self.purchases_df, self.products_df, on="商品编号")
            merged['消费金额'] = merged['数量'] * merged['单价'] * (merged['折扣率'] / 100)
            result = merged.groupby('客户编号')['消费金额'].sum().reset_index()
            result = result.sort_values('消费金额', ascending=False)
            self.print_calculation_formulas(merged)
            self.display_results(result)
            # 保存明文计算结果
            self.plaintext_results = result
            self.update_status("明文计算完成")
        except Exception as e:
            messagebox.showerror("错误", f"明文计算失败: {str(e)}")

    def print_calculation_formulas(self, merged_df):
        customer_groups = merged_df.groupby('客户编号')
        for customer_id, group in customer_groups:
            print(f"客户编号 {customer_id} 的消费计算式：")
            formulas = []
            for _, row in group.iterrows():
                formula = f"({row['数量']} * {row['单价']} * ({row['折扣率']}/100))"
                formulas.append(formula)
            total_formula = " + ".join(formulas)
            print(total_formula)
            print(f"总和为：{eval(total_formula)}")
            print("-" * 50)

    def analyze_encrypted(self):
        if self.products_df is None or self.purchases_df is None:
            messagebox.showerror("错误", "请先加载商品数据和购买记录")
            return
        try:
            self.update_status("正在加密数据并计算...")
            merged = pd.merge(self.purchases_df, self.products_df, on="商品编号")
            customer_groups = merged.groupby('客户编号')
            results = []
            for customer_id, group in customer_groups:
                quantities = group['数量'].values
                prices = group['单价'].values
                discounts = (group['折扣率'].values / 100)
                encrypted_total = self.fhe.encrypt([0.0])
                for q, p, d in zip(quantities, prices, discounts):
                    encrypted_q = self.fhe.encrypt([q])
                    encrypted_p = self.fhe.encrypt([p])
                    encrypted_d = self.fhe.encrypt([d])
                    encrypted_product = self.fhe.multiply(encrypted_q, encrypted_p)
                    encrypted_product = self.fhe.multiply(encrypted_product, encrypted_d)
                    encrypted_total = self.fhe.add(encrypted_total, encrypted_product)
                total_spent = self.fhe.decrypt(encrypted_total)[0]
                results.append({'客户编号': customer_id, '消费金额': total_spent})
            result = pd.DataFrame(results).sort_values('消费金额', ascending=False)
            self.display_results(result)
            # 保存加密计算结果
            self.encrypted_results = result
            self.update_status("加密计算完成")
        except Exception as e:
            messagebox.showerror("错误", f"加密计算失败: {str(e)}")

    def display_results(self, result_df):
        for item in self.tree.get_children():
            self.tree.delete(item)
        for _, row in result_df.iterrows():
            self.tree.insert("", "end", values=(row['客户编号'], f"{row['消费金额']:.2f}"))

    def visualize_results(self):
        # 检查明文和加密结果是否都存在
        if self.plaintext_results is None or self.encrypted_results is None:
            messagebox.showwarning("警告", "请先完成明文和加密计算")
            return
        # 提取明文数据
        plain_ids = self.plaintext_results['客户编号'].astype(str).tolist()
        plain_amounts = self.plaintext_results['消费金额'].tolist()

        # 提取加密数据
        encrypted_ids = self.encrypted_results['客户编号'].astype(str).tolist()
        encrypted_amounts = self.encrypted_results['消费金额'].tolist()

        # 确保客户编号一致（按明文排序对齐）
        if plain_ids != encrypted_ids:
            # 若加密结果顺序不同，重新对齐（以明文为基准）
            encrypted_amounts_aligned = []
            for cid in plain_ids:
                # 找到加密结果中对应客户编号的金额
                idx = self.encrypted_results[self.encrypted_results['客户编号'] == cid].index[0]
                encrypted_amounts_aligned.append(self.encrypted_results.iloc[idx]['消费金额'])
            encrypted_amounts = encrypted_amounts_aligned

        plt.rcParams['font.sans-serif'] = ['SimHei']  # 中文支持
        plt.rcParams['axes.unicode_minus'] = False   # 负号显示
        plt.figure(figsize=(10, 6))

        # 绘制明文折线（蓝色）
        plt.plot(plain_ids, plain_amounts, label='明文计算', color='blue', marker='o', linestyle='-', linewidth=2)
        # 绘制加密折线（橙色）
        plt.plot(plain_ids, encrypted_amounts, label='加密计算', color='orange', marker='s', linestyle='--', linewidth=2)

        # 添加数值标注
        for x, y in zip(plain_ids, plain_amounts):
            plt.text(x, y + 500, f'{y:.2f}', ha='center', va='bottom', fontsize=18, color='blue')
        for x, y in zip(plain_ids, encrypted_amounts):
            plt.text(x, y - 1000, f'{y:.2f}', ha='center', va='top', fontsize=18, color='orange')

        plt.title('客户消费金额对比（明文 vs 加密）')
        plt.xlabel('客户编号')
        plt.ylabel('消费金额')
        plt.xticks(rotation=45)
        plt.legend()
        plt.grid(alpha=0.3)  # 添加网格辅助观察
        plt.tight_layout()
        plt.show()


class FHEEncryption:
    def __init__(self, scheme: str = "CKKS", poly_modulus_degree: int = 8192,
                 coeff_mod_bit_sizes: List[int] = [60, 40, 40, 60],
                 scale: int = 2 ** 40, plain_modulus: int = 65537):
        self.scheme = scheme
        self.poly_modulus_degree = poly_modulus_degree
        self.coeff_mod_bit_sizes = coeff_mod_bit_sizes
        self.scale = scale
        self.plain_modulus = plain_modulus
        self.context = self._create_context()
        self.context.generate_galois_keys()
        self.context.generate_relin_keys()

    def _create_context(self) -> ts.Context:
        if self.scheme == "CKKS":
            context = ts.context(
                ts.SCHEME_TYPE.CKKS,
                poly_modulus_degree=self.poly_modulus_degree,
                coeff_mod_bit_sizes=self.coeff_mod_bit_sizes
            )
            # 设置全局缩放因子（修复no global scale错误）
            context.global_scale = self.scale
            return context
        elif self.scheme == "BFV":
            context = ts.context(
                ts.SCHEME_TYPE.BFV,
                poly_modulus_degree=self.poly_modulus_degree,
                plain_modulus=self.plain_modulus
            )
            return context
        else:
            raise ValueError(f"不支持的加密方案: {self.scheme}")

    def encrypt(self, data: Union[List[float], List[int], np.ndarray]) -> Union[ts.CKKSTensor, ts.BFVTensor]:
        if isinstance(data, np.ndarray):
            data = data.tolist()
        if self.scheme == "CKKS":
            return ts.ckks_vector(self.context, data)
        else:
            return ts.bfv_vector(self.context, data)

    def decrypt(self, encrypted_data: Union[ts.CKKSTensor, ts.BFVTensor]) -> List[Union[float, int]]:
        return encrypted_data.decrypt()

    def add(self, encrypted_data1: Union[ts.CKKSTensor, ts.BFVTensor],
            encrypted_data2: Union[ts.CKKSTensor, ts.BFVTensor]) -> Union[ts.CKKSTensor, ts.BFVTensor]:
        return encrypted_data1 + encrypted_data2

    def subtract(self, encrypted_data1: Union[ts.CKKSTensor, ts.BFVTensor],
                 encrypted_data2: Union[ts.CKKSTensor, ts.BFVTensor]) -> Union[ts.CKKSTensor, ts.BFVTensor]:
        return encrypted_data1 - encrypted_data2

    def multiply(self, encrypted_data1: Union[ts.CKKSTensor, ts.BFVTensor],
                 encrypted_data2: Union[ts.CKKSTensor, ts.BFVTensor]) -> Union[ts.CKKSTensor, ts.BFVTensor]:
        return encrypted_data1 * encrypted_data2

    def dot(self, encrypted_data1: Union[ts.CKKSTensor, ts.BFVTensor],
            encrypted_data2: Union[ts.CKKSTensor, ts.BFVTensor]) -> Union[ts.CKKSTensor, ts.BFVTensor]:
        return encrypted_data1.dot(encrypted_data2)

    def matmul(self, encrypted_matrix: ts.CKKSTensor,
               matrix: Union[List[List[float]], np.ndarray]) -> ts.CKKSTensor:
        if self.scheme != "CKKS":
            raise ValueError("矩阵乘法仅支持CKKS方案")
        if isinstance(matrix, np.ndarray):
            matrix = matrix.tolist()
        vector_size = len(encrypted_matrix.decrypt())
        matrix_rows = len(matrix)
        matrix_cols = len(matrix[0]) if matrix_rows > 0 else 0
        if vector_size != matrix_rows:
            raise ValueError(f"矩阵形状不匹配：向量大小为{vector_size}，但矩阵有{matrix_rows}行")
        return encrypted_matrix.mm(matrix)

    def polynomial(self, encrypted_data: ts.CKKSTensor,
                   coefficients: List[float]) -> ts.CKKSTensor:
        if self.scheme != "CKKS":
            raise ValueError("多项式计算仅支持CKKS方案")
        result = self.encrypt([coefficients[-1]])
        for c in reversed(coefficients[:-1]):
            result = self.add(self.multiply(result, encrypted_data), self.encrypt([c]))
        return result

    def save_context(self, path: str):
        with open(path, "wb") as f:
            f.write(self.context.serialize())

    def load_context(self, path: str):
        with open(path, "rb") as f:
            self.context = ts.context_from(f.read())


if __name__ == "__main__":
    root = Tk()
    app = PrivacyPreservingSpendingAnalyzer(root)
    root.mainloop()